#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   train.py
@Time    :   2021/11/09 15:24:41
@Author  :   Yaadon 
'''

# here put the import lib
import paddle
import paddle.nn.functional as F
from net.LSTM import SentimentClassifier
from data.loader import load_imdb,data_preprocess,build_batch, build_dict, convert_corpus_to_id

# 定义训练参数
epoch_num = 10
batch_size = 128

learning_rate = 0.01
dropout_rate = 0.2
num_layers = 1
hidden_size = 256
embedding_size = 256
max_seq_len = 128

train_corpus = load_imdb(True)
# test_corpus = load_imdb(False)
train_corpus = data_preprocess(train_corpus)
# test_corpus = data_preprocess(test_corpus)
# print(train_corpus[:5])
# print(test_corpus[:5])
word2id_freq, word2id_dict = build_dict(train_corpus)
train_corpus = convert_corpus_to_id(train_corpus, word2id_dict)
# test_corpus = convert_corpus_to_id(test_corpus, word2id_dict)
word2id_freq, word2id_dict = build_dict(train_corpus)
vocab_size = len(word2id_freq)

# 检测是否可以使用GPU，如果可以优先使用GPU
use_gpu = True if paddle.get_device().startswith("gpu") else False
if use_gpu:
    paddle.set_device('gpu:0')

# 实例化模型
sentiment_classifier = SentimentClassifier(hidden_size, vocab_size, embedding_size,  num_steps=max_seq_len, num_layers=num_layers, dropout_rate=dropout_rate)

# 指定优化策略，更新模型参数
optimizer = paddle.optimizer.Adam(learning_rate=learning_rate, beta1=0.9, beta2=0.999, parameters= sentiment_classifier.parameters()) 

# 定义训练函数
# 记录训练过程中的损失变化情况，可用于后续画图查看训练情况
losses = []
steps = []

def train(model):
    # 开启模型训练模式
    model.train()
    
    # 建立训练数据生成器，每次迭代生成一个batch，每个batch包含训练文本和文本对应的情感标签
    train_loader = build_batch(word2id_dict, train_corpus, batch_size, epoch_num, max_seq_len)
    
    for step, (sentences, labels) in enumerate(train_loader):
        # 获取数据，并将张量转换为Tensor类型
        sentences = paddle.to_tensor(sentences)
        labels = paddle.to_tensor(labels)
        
        # 前向计算，将数据feed进模型，并得到预测的情感标签和损失
        logits = model(sentences)

        # 计算损失
        loss = F.cross_entropy(input=logits, label=labels, soft_label=False)
        loss = paddle.mean(loss)

        # 后向传播
        loss.backward()
        # 更新参数
        optimizer.step()
        # 清除梯度
        optimizer.clear_grad()

        if step % 100 == 0:
            # 记录当前步骤的loss变化情况
            losses.append(loss.numpy()[0])
            steps.append(step)
            # 打印当前loss数值
            print("step %d, loss %.3f" % (step, loss.numpy()[0]))


if __name__ == '__main__':
    #训练模型
    train(sentiment_classifier)

    # 保存模型，包含两部分：模型参数和优化器参数
    model_name = "sentiment_classifier"
    # 保存训练好的模型参数
    paddle.save(sentiment_classifier.state_dict(), "{}.pdparams".format(model_name))
    # 保存优化器参数，方便后续模型继续训练
    paddle.save(optimizer.state_dict(), "{}.pdopt".format(model_name))
